// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cassert>
#include <cstddef>
#include <iterator>
#include <optional>
#include <utility>

#include <nebula/core/chunked_array.h>
#include <nebula/types/type.h>
#include <nebula/types/type_fwd.h>
#include <nebula/types/type_traits.h>
#include <turbo/base/macros.h>

namespace nebula {
    namespace stl {

        namespace detail {

            template<typename ArrayType>
            struct DefaultValueAccessor {
                using value_type = decltype(std::declval<ArrayType>().get_view(0));

                value_type operator()(const ArrayType &array, int64_t index) {
                    return array.get_view(index);
                }
            };

        }  // namespace detail

        template<typename ArrayType,
                typename ValueAccessor = detail::DefaultValueAccessor<ArrayType>>
        class ArrayIterator {
        public:
            using value_type = std::optional<typename ValueAccessor::value_type>;
            using difference_type = int64_t;
            using pointer = value_type *;
            using reference = value_type &;
            using iterator_category = std::random_access_iterator_tag;

            // Some algorithms need to default-construct an iterator
            ArrayIterator() : array_(nullptr), index_(0) {}

            explicit ArrayIterator(const ArrayType &array, int64_t index = 0)
                    : array_(&array), index_(index) {}

            // Value access
            value_type operator*() const {
                assert(array_);
                return array_->is_null(index_) ? value_type{} : array_->get_view(index_);
            }

            value_type operator[](difference_type n) const {
                assert(array_);
                return array_->is_null(index_ + n) ? value_type{} : array_->get_view(index_ + n);
            }

            int64_t index() const { return index_; }

            // Forward / backward
            ArrayIterator &operator++() {
                ++index_;
                return *this;
            }

            ArrayIterator &operator--() {
                --index_;
                return *this;
            }

            ArrayIterator operator++(int) {
                ArrayIterator tmp(*this);
                ++index_;
                return tmp;
            }

            ArrayIterator operator--(int) {
                ArrayIterator tmp(*this);
                --index_;
                return tmp;
            }

            // Arithmetic
            difference_type operator-(const ArrayIterator &other) const {
                return index_ - other.index_;
            }

            ArrayIterator operator+(difference_type n) const {
                return ArrayIterator(*array_, index_ + n);
            }

            ArrayIterator operator-(difference_type n) const {
                return ArrayIterator(*array_, index_ - n);
            }

            friend inline ArrayIterator operator+(difference_type diff,
                                                  const ArrayIterator &other) {
                return ArrayIterator(*other.array_, diff + other.index_);
            }

            friend inline ArrayIterator operator-(difference_type diff,
                                                  const ArrayIterator &other) {
                return ArrayIterator(*other.array_, diff - other.index_);
            }

            ArrayIterator &operator+=(difference_type n) {
                index_ += n;
                return *this;
            }

            ArrayIterator &operator-=(difference_type n) {
                index_ -= n;
                return *this;
            }

            // Comparisons
            bool operator==(const ArrayIterator &other) const { return index_ == other.index_; }

            bool operator!=(const ArrayIterator &other) const { return index_ != other.index_; }

            bool operator<(const ArrayIterator &other) const { return index_ < other.index_; }

            bool operator>(const ArrayIterator &other) const { return index_ > other.index_; }

            bool operator<=(const ArrayIterator &other) const { return index_ <= other.index_; }

            bool operator>=(const ArrayIterator &other) const { return index_ >= other.index_; }

        private:
            const ArrayType *array_;
            int64_t index_;
        };

        template<typename ArrayType,
                typename ValueAccessor = detail::DefaultValueAccessor<ArrayType>>
        class ChunkedArrayIterator {
        public:
            using value_type = std::optional<typename ValueAccessor::value_type>;
            using difference_type = int64_t;
            using pointer = value_type *;
            using reference = value_type &;
            using iterator_category = std::random_access_iterator_tag;

            // Some algorithms need to default-construct an iterator
            ChunkedArrayIterator() noexcept: chunked_array_(nullptr), index_(0) {}

            explicit ChunkedArrayIterator(const ChunkedArray &chunked_array,
                                          int64_t index = 0) noexcept
                    : chunked_array_(&chunked_array), index_(index) {}

            // Value access
            value_type operator*() const {
                auto chunk_location = GetChunkLocation(index_);
                ArrayIterator<ArrayType> target_iterator{
                        turbo::checked_cast<const ArrayType &>(
                                *chunked_array_->chunk(static_cast<int>(chunk_location.chunk_index)))};
                return target_iterator[chunk_location.index_in_chunk];
            }

            value_type operator[](difference_type n) const { return *(*this + n); }

            int64_t index() const { return index_; }

            // Forward / backward
            ChunkedArrayIterator &operator++() {
                (*this) += 1;
                return *this;
            }

            ChunkedArrayIterator &operator--() {
                (*this) -= 1;
                return *this;
            }

            ChunkedArrayIterator operator++(int) {
                ChunkedArrayIterator tmp(*this);
                ++*this;
                return tmp;
            }

            ChunkedArrayIterator operator--(int) {
                ChunkedArrayIterator tmp(*this);
                --*this;
                return tmp;
            }

            // Arithmetic
            difference_type operator-(const ChunkedArrayIterator &other) const {
                return index_ - other.index_;
            }

            ChunkedArrayIterator operator+(difference_type n) const {
                assert(chunked_array_);
                return ChunkedArrayIterator(*chunked_array_, index_ + n);
            }

            ChunkedArrayIterator operator-(difference_type n) const {
                assert(chunked_array_);
                return ChunkedArrayIterator(*chunked_array_, index_ - n);
            }

            friend inline ChunkedArrayIterator operator+(difference_type diff,
                                                         const ChunkedArrayIterator &other) {
                assert(other.chunked_array_);
                return ChunkedArrayIterator(*other.chunked_array_, diff + other.index_);
            }

            friend inline ChunkedArrayIterator operator-(difference_type diff,
                                                         const ChunkedArrayIterator &other) {
                assert(other.chunked_array_);
                return ChunkedArrayIterator(*other.chunked_array_, diff - other.index_);
            }

            ChunkedArrayIterator &operator+=(difference_type n) {
                index_ += n;
                return *this;
            }

            ChunkedArrayIterator &operator-=(difference_type n) {
                (*this) += -n;
                return *this;
            }

            // Comparisons
            bool operator==(const ChunkedArrayIterator &other) const {
                return index_ == other.index_;
            }

            bool operator!=(const ChunkedArrayIterator &other) const {
                return index_ != other.index_;
            }

            bool operator<(const ChunkedArrayIterator &other) const {
                return index_ < other.index_;
            }

            bool operator>(const ChunkedArrayIterator &other) const {
                return index_ > other.index_;
            }

            bool operator<=(const ChunkedArrayIterator &other) const {
                return index_ <= other.index_;
            }

            bool operator>=(const ChunkedArrayIterator &other) const {
                return index_ >= other.index_;
            }

        private:
            nebula::internal::ChunkLocation GetChunkLocation(int64_t index) const {
                assert(chunked_array_);
                return chunked_array_->chunk_resolver_.resolve(index);
            }

            const ChunkedArray *chunked_array_;
            int64_t index_;
        };

        /// Return an iterator to the beginning of the chunked array
        template<typename Type, typename ArrayType = typename TypeTraits<Type>::ArrayType>
        ChunkedArrayIterator<ArrayType> begin(const ChunkedArray &chunked_array) {
            return ChunkedArrayIterator<ArrayType>(chunked_array);
        }

        /// Return an iterator to the end of the chunked array
        template<typename Type, typename ArrayType = typename TypeTraits<Type>::ArrayType>
        ChunkedArrayIterator<ArrayType> end(const ChunkedArray &chunked_array) {
            return ChunkedArrayIterator<ArrayType>(chunked_array, chunked_array.length());
        }

        template<typename ArrayType>
        struct ChunkedArrayRange {
            const ChunkedArray *chunked_array;

            ChunkedArrayIterator<ArrayType> begin() {
                return stl::ChunkedArrayIterator<ArrayType>(*chunked_array);
            }

            ChunkedArrayIterator<ArrayType> end() {
                return stl::ChunkedArrayIterator<ArrayType>(*chunked_array, chunked_array->length());
            }
        };

        /// Return an iterable range over the chunked array
        template<typename Type, typename ArrayType = typename TypeTraits<Type>::ArrayType>
        ChunkedArrayRange<ArrayType> Iterate(const ChunkedArray &chunked_array) {
            return stl::ChunkedArrayRange<ArrayType>{&chunked_array};
        }

    }  // namespace stl
}  // namespace nebula

namespace std {

    template<typename ArrayType>
    struct iterator_traits<::nebula::stl::ArrayIterator<ArrayType>> {
        using IteratorType = ::nebula::stl::ArrayIterator<ArrayType>;
        using difference_type = typename IteratorType::difference_type;
        using value_type = typename IteratorType::value_type;
        using pointer = typename IteratorType::pointer;
        using reference = typename IteratorType::reference;
        using iterator_category = typename IteratorType::iterator_category;
    };

    template<typename ArrayType>
    struct iterator_traits<::nebula::stl::ChunkedArrayIterator<ArrayType>> {
        using IteratorType = ::nebula::stl::ChunkedArrayIterator<ArrayType>;
        using difference_type = typename IteratorType::difference_type;
        using value_type = typename IteratorType::value_type;
        using pointer = typename IteratorType::pointer;
        using reference = typename IteratorType::reference;
        using iterator_category = typename IteratorType::iterator_category;
    };

}  // namespace std
